import matplotlib.patches as patches
import matplotlib.pyplot as plt
import random

from robot_env import CollisionManager, PathPlanner, Map, TerraNodes, MovingAgent
from walle_world import WalleWorld

class MyMovingAgent(MovingAgent):
    def __init__(self, agent_id, current_node=None, z_index=0):
        MovingAgent.__init__(self, agent_id, current_node, z_index)

    def calc_moving_time(self, next_terra):
        return 1

    def __str__(self):
        return self.agent_id


class MyCollisionManager(CollisionManager):
    def __init__(self, agent_id):
        CollisionManager.__init__(self, agent_id)

    def on_request_move_to_terra_event(self, tick, environment, agents, event):
        # to determine if collision will happen
        yield self.build_event_instance(tick, tick, "approved_move_to_terra_event", event.message, event.source)

    def __str__(self):
        return self.agent_id

class MyPathPlanner(PathPlanner):
    def __init__(self, agent_id):
        PathPlanner.__init__(self, agent_id)

    def __str__(self):
        return self.agent_id

    def get_next_node(self):
        return str(random.randint(1, 49))

    def on_walle_init_event(self, tick, environment, agents, event):
        new_event = {"target": environment["map"].get_node_by_id(self.get_next_node())}
        yield self.build_event_instance(tick, tick + 1, "request_move_to_target_event", new_event, "moving_agent_1")

    def on_request_moving_plan_event(self, tick, environment, agents, event):
        agent = event.message["agent"]
        target = event.message["moving_target"]
        start_id = int(agent.current_node.node_id)
        end_id = int(target.node_id)
        plan = [agent.current_node]
        while start_id != end_id:
            if start_id == 34:
                node = environment["map"].get_node_by_id(str(20))
                plan.append(node)
            node = environment["map"].get_node_by_id(str(start_id))
            plan.append(node)
            start_id = start_id + 1
            if start_id > 49:
                start_id = 1
        plan.append(target)
        yield self.build_event_instance(tick, tick, "moving_plan_event", {"moving_plan": plan}, event.source)

    def on_moved_to_target_event(self, tick, environment, agents, event):
        new_event = {"target": environment["map"].get_node_by_id(self.get_next_node())}
        yield self.build_event_instance(tick, tick + 1, "request_move_to_target_event", new_event, event.source)


if __name__ == "__main__":
    nodes = []
    nodes.append(TerraNodes("1", 0, 0, [0,0,0,1,0,1,0,0]))
    nodes.append(TerraNodes("2", 1, 0, [0,0,0,1,0,0,0,1]))
    nodes.append(TerraNodes("3", 2, 0, [0,0,0,1,0,0,0,1]))
    nodes.append(TerraNodes("4", 3, 0, [0,0,0,1,0,0,0,1]))
    nodes.append(TerraNodes("5", 4, 0, [0,0,0,1,0,0,0,1]))
    nodes.append(TerraNodes("6", 5, 0, [0,0,0,1,0,0,0,1]))
    nodes.append(TerraNodes("7", 6, 0, [0,0,0,1,0,0,0,1]))
    nodes.append(TerraNodes("8", 7, 0, [0,0,0,1,0,0,0,1]))
    nodes.append(TerraNodes("9", 8, 0, [0,0,0,1,0,0,0,1]))
    nodes.append(TerraNodes("10", 9, 0, [0,0,0,0,0,1,0,1]))
    nodes.append(TerraNodes("11", 9, 1, [0,1,0,0,0,1,0,0]))
    nodes.append(TerraNodes("12", 9, 2, [0,1,0,0,0,0,0,1]))
    nodes.append(TerraNodes("13", 8, 2, [0,0,0,1,0,0,0,1]))
    nodes.append(TerraNodes("14", 7, 2, [0,0,0,1,0,0,0,1]))
    nodes.append(TerraNodes("15", 6, 2, [0,0,0,1,0,0,0,1]))
    nodes.append(TerraNodes("16", 5, 2, [0,0,0,1,0,1,0,0]))
    nodes.append(TerraNodes("17", 5, 3, [0,1,0,0,0,1,0,0]))
    nodes.append(TerraNodes("18", 5, 4, [0,1,0,0,0,1,0,0]))
    nodes.append(TerraNodes("19", 5, 5, [0,1,0,0,0,1,0,0]))
    nodes.append(TerraNodes("20", 5, 6, [0,1,0,1,0,1,0,1]))
    nodes.append(TerraNodes("21", 6, 6, [0,0,0,1,0,0,0,1]))
    nodes.append(TerraNodes("22", 7, 6, [0,0,0,1,0,0,0,1]))
    nodes.append(TerraNodes("23", 8, 6, [0,0,0,1,0,0,0,1]))
    nodes.append(TerraNodes("24", 9, 6, [0,0,0,0,0,1,0,1]))
    nodes.append(TerraNodes("25", 9, 7, [0,1,0,0,0,1,0,0]))
    nodes.append(TerraNodes("26", 9, 8, [0,1,0,0,0,1,0,0]))
    nodes.append(TerraNodes("27", 9, 9, [0,1,0,0,0,0,0,1]))
    nodes.append(TerraNodes("28", 8, 9, [0,0,0,1,0,0,0,1]))
    nodes.append(TerraNodes("29", 7, 9, [0,0,0,1,0,0,0,1]))
    nodes.append(TerraNodes("30", 6, 9, [0,0,0,1,0,0,0,1]))
    nodes.append(TerraNodes("31", 5, 9, [0,1,0,1,0,0,0,0]))
    nodes.append(TerraNodes("32", 5, 8, [0,1,0,0,0,1,0,0]))
    nodes.append(TerraNodes("33", 5, 7, [0,1,0,0,0,1,0,0]))
    nodes.append(TerraNodes("34", 4, 6, [0,0,0,1,0,0,0,1]))
    nodes.append(TerraNodes("35", 3, 6, [0,0,0,1,0,0,0,1]))
    nodes.append(TerraNodes("36", 2, 6, [0,0,0,1,0,1,0,0]))
    nodes.append(TerraNodes("37", 2, 7, [0,1,0,0,0,1,0,0]))
    nodes.append(TerraNodes("38", 2, 8, [0,1,0,0,0,1,0,0]))
    nodes.append(TerraNodes("39", 2, 9, [0,1,0,0,0,0,0,1]))
    nodes.append(TerraNodes("40", 1, 9, [0,0,0,1,0,0,0,1]))
    nodes.append(TerraNodes("41", 0, 9, [0,1,0,1,0,0,0,0]))
    nodes.append(TerraNodes("42", 0, 8, [0,1,0,0,0,1,0,0]))
    nodes.append(TerraNodes("43", 0, 7, [0,1,0,0,0,1,0,0]))
    nodes.append(TerraNodes("44", 0, 6, [0,1,0,0,0,1,0,0]))
    nodes.append(TerraNodes("45", 0, 5, [0,1,0,0,0,1,0,0]))
    nodes.append(TerraNodes("46", 0, 4, [0,1,0,0,0,1,0,0]))
    nodes.append(TerraNodes("47", 0, 3, [0,1,0,0,0,1,0,0]))
    nodes.append(TerraNodes("48", 0, 2, [0,1,0,0,0,1,0,0]))
    nodes.append(TerraNodes("49", 0, 1, [0,1,0,0,0,1,0,0]))

    map = Map("test_map", 10, 10, nodes)

    walle_world = WalleWorld("test_robot_world", {"map": map})

    # render options
    terra_width = 1
    terra_height = 1
    figure = plt.figure()
    ax = figure.add_subplot(1, 1, 1, aspect='equal')
    ax.set_xlim(0, map.width * terra_width)
    ax.set_ylim(0, map.height * terra_width)
    ax.invert_yaxis()
    def render(tick, environment, agents):
        if tick % 1 == 0:
            ax.cla()
            #draw terra
            for terra in environment["map"].nodes:
                x = terra.x * terra_width
                y = terra.y * terra_width
                width = terra_width
                height = terra_width
                ax.add_patch(patches.Rectangle((x, y), width, height, fill=True, facecolor="#f1f1f1",
                                               edgecolor='gray', linewidth=0.5, linestyle="solid"))
            #draw robot
            for a in agents:
                if a["agent_id"] == "moving_agent_1":
                    agent = a["body"]
                    agent_x = agent.x * terra_width + 0.2 * terra_width
                    agent_y = agent.y * terra_height + 0.2 * terra_height
                    width = 0.6 * terra_width
                    height = 0.6 * terra_height
                    ax.add_patch(patches.Rectangle((agent_x, agent_y), width, height, fill=True, facecolor="red",
                                                   edgecolor='gray', linewidth=0.5, linestyle="solid"))
            figure.show()
            plt.pause(0.1)

    walle_world.add_world_observer(render)

    agent0 = MyCollisionManager("collision_manager")
    agent1 = MyPathPlanner("path_planner")
    agent2 = MyMovingAgent("moving_agent_1", map.get_node_by_id("1"))
    walle_world.add_agent(agent0)
    walle_world.add_agent(agent1)
    walle_world.add_agent(agent2)

    #start running
    walle_world.run(10000)

    walle_world.dump_event_log()

    while True:
        plt.pause(1.0)
